#!/usr/bin/python3

import sys
import os
import PIL.Image
from PIL import Image
import torch
import numpy as np
import cv2
import onnxruntime
from utils import snapshot_download
from torch import multiprocessing
from preprocessing import Blipv2
from modelscope.pipelines import pipeline


def skin_retouching(tmp_path):
    print(tmp_path)
    output_path = os.path.dirname(tmp_path)
    output_name = os.path.basename(tmp_path)
    output_path_name = os.path.join(output_path + "/skin", 'skin_' + output_name)
    skin_retouching = pipeline('skin-retouching-torch', model='damo/cv_unet_skin_retouching_torch',
                               model_revision='v1.0.1')
    result = skin_retouching(tmp_path)
    if (result is None or (result['output_img'] is None)):
        print('Cannot do skin retouching, do not use this image.')
    cv2.imwrite(output_path_name, result['output_img'])


def main(argv):
    format_data_dir = 'E:/data/output'
    img_list = ['E:/data/test/test_20231124104818.jpg']
    skin_path = 'E:/data/test/test_20231124104715.jpg'
    # skin_retouching(skin_path)
    # prepare_dataset(img_list, format_data_dir)  # 检测图形是否旋转
    # data_process_fn(format_data_dir, True)  # 抠图、生成标签
    skin_dir_path = 'E:/data/test/'
    files = os.listdir(skin_dir_path)
    for file in files:
        file_path = skin_dir_path + file
        if os.path.isdir(file_path):
            continue
        skin_retouching(file_path)


def _data_process_fn_process(input_img_dir):
    Blipv2()(input_img_dir)


def data_process_fn(input_img_dir, use_data_process):
    ## TODO add face quality filter
    if use_data_process:
        ## TODO
        _data_process_fn_process(input_img_dir)
        # _process = multiprocessing.Process(target=_data_process_fn_process, args=(input_img_dir,))
        # _process.start()
        # _process.join()

    return os.path.join(str(input_img_dir) + '_labeled', "metadata.jsonl")


def prepare_dataset(instance_images: list, output_dataset_dir):
    if not os.path.exists(output_dataset_dir):
        os.makedirs(output_dataset_dir)
    for i, temp_path in enumerate(instance_images):
        image = PIL.Image.open(temp_path)
        # image = PIL.Image.open(temp_path.name)
        '''
        w, h = image.size
        max_size = max(w, h)
        ratio =  1024 / max_size
        new_w = round(w * ratio)
        new_h = round(h * ratio)
        '''
        image = image.convert('RGB')
        image = get_rot(image)
        # image = image.resize((new_w, new_h))
        # image = image.resize((new_w, new_h), PIL.Image.ANTIALIAS)
        out_path = f'{output_dataset_dir}/{i:03d}.jpg'
        image.save(out_path, format='JPEG', quality=100)


def get_rot(image):
    model_dir = snapshot_download('Cherrytest/rot_bgr',
                                  revision='v1.0.0')
    model_path = os.path.join(model_dir, 'rot_bgr.onnx')
    providers = ['CPUExecutionProvider']
    if torch.cuda.is_available():
        providers.insert(0, 'CUDAExecutionProvider')
    ort_session = onnxruntime.InferenceSession(model_path, providers=providers)

    img_cv = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
    img_clone = img_cv.copy()
    img_np = cv2.resize(img_cv, (224, 224))
    img_np = img_np.astype(np.float32)
    mean = np.array([103.53, 116.28, 123.675], dtype=np.float32).reshape((1, 1, 3))
    norm = np.array([0.01742919, 0.017507, 0.01712475], dtype=np.float32).reshape((1, 1, 3))
    img_np = (img_np - mean) * norm
    img_tensor = torch.from_numpy(img_np)
    img_tensor = img_tensor.unsqueeze(0)
    img_nchw = img_tensor.permute(0, 3, 1, 2)
    ort_inputs = {ort_session.get_inputs()[0].name: img_nchw.numpy()}
    outputs = ort_session.run(None, ort_inputs)
    logits = outputs[0].reshape((-1,))
    probs = softmax(logits)
    rot_idx = np.argmax(probs)
    if rot_idx == 1:
        print('rot 90')
        img_clone = cv2.transpose(img_clone)
        img_clone = np.flip(img_clone, 1)
        return Image.fromarray(cv2.cvtColor(img_clone, cv2.COLOR_BGR2RGB))
    elif rot_idx == 2:
        print('rot 180')
        img_clone = cv2.flip(img_clone, -1)
        return Image.fromarray(cv2.cvtColor(img_clone, cv2.COLOR_BGR2RGB))
    elif rot_idx == 3:
        print('rot 270')
        img_clone = cv2.transpose(img_clone)
        img_clone = np.flip(img_clone, 0)
        return Image.fromarray(cv2.cvtColor(img_clone, cv2.COLOR_BGR2RGB))
    else:
        return image


def softmax(x):
    x -= np.max(x, axis=0, keepdims=True)
    x = np.exp(x) / np.sum(np.exp(x), axis=0, keepdims=True)
    return x


if __name__ == "__main__":
    main(sys.argv[1:])
